import asyncio
import logging
from typing import Dict, Set

from app.database.redis import get_redis
from app.schemas.stomp import StompCommand, StompFrame

logger = logging.getLogger(__name__)


class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, asyncio.Queue] = {}  # user_id -> send_queue
        self.subscriptions: Dict[str, Set[str]] = {}  # user_id -> set of topics
        self.user_sockets: Dict[str, str] = {}  # socket_id -> user_id
        self.next_id = 0

    async def connect(self, websocket, user_id: str):
        socket_id = f"sock_{self.next_id}"
        self.next_id += 1
        self.user_sockets[socket_id] = user_id
        self.active_connections[user_id] = asyncio.Queue()
        self.subscriptions[user_id] = set()
        logger.info(f"User {user_id} connected (socket: {socket_id})")
        asyncio.create_task(self._send_loop(websocket, user_id))

    async def disconnect(self, websocket, user_id: str):
        if user_id in self.active_connections:
            del self.active_connections[user_id]
        if user_id in self.subscriptions:
            del self.subscriptions[user_id]
        # Remove socket mapping
        for sock_id, uid in list(self.user_sockets.items()):
            if uid == user_id:
                del self.user_sockets[sock_id]
        logger.info(f"User {user_id} disconnected")

    async def _send_loop(self, websocket, user_id: str):
        try:
            while True:
                message = await self.active_connections[user_id].get()
                await websocket.send_text(message)
        except Exception as e:
            logger.error(f"Send loop error for {user_id}: {e}")

    async def subscribe(self, user_id: str, topic: str):
        self.subscriptions[user_id].add(topic)
        logger.info(f"User {user_id} subscribed to {topic}")

    async def unsubscribe(self, user_id: str, topic: str):
        self.subscriptions[user_id].discard(topic)

    async def broadcast_to_topic(self, topic: str, message: str):
        """通过 Redis 广播，由后台任务消费后推送给订阅者"""
        redis = await get_redis()
        await redis.publish(f"stomp:topic:{topic}", message)

    async def send_to_user(self, user_id: str, message: str):
        if user_id in self.active_connections:
            await self.active_connections[user_id].put(message)

    async def handle_redis_message(self, channel: str, data: str):
        """从 Redis 接收广播消息并推送给订阅者"""
        if channel.startswith("stomp:topic:"):
            topic = channel.replace("stomp:topic:", "", 1)
            frame = StompFrame(
                command=StompCommand.SEND, headers={"destination": f"/topic/{topic}"}, body=data
            ).to_frame()
            # 推送给所有订阅该主题的用户
            for uid, topics in self.subscriptions.items():
                if topic in topics:
                    await self.send_to_user(uid, frame)


manager = ConnectionManager()
